#include "scheduler.h"
#include "interval_tree.h"
#include "types.h"
#include <cstdlib>
#include <map>
#include <algorithm>
#include <memory>
#include <vector>
using namespace std;

extern IntervalTree * globalMemTree;

// 构造函数：初始化调度器状态
Scheduler::Scheduler(long long m, const std::vector<Access> & accesses) : memoryCapacity(m), accessVec(accesses), currentTime(0)
{
    LOG("accesses: %d", accesses.size());
    groupAccessesIntoTasks();
    globalMemTree = &memoryTree;
}

// 将访问请求按 `start_time` 分组到任务中
void Scheduler::groupAccessesIntoTasks()
{
    std::map<unsigned long long, vector<Access *>> accessGroupMap;
    LOG("accessVec: %d", accessVec.size());
    for (auto & access : accessVec) {
        LOG("time %lld access %p", access.timeStart, &access);
        accessGroupMap[access.timeStart].push_back(&access);
    }
    for (const auto & group : accessGroupMap) {
        shared_ptr<Task> task(new Task);
        LOG("group %lld mem %d", group.first, group.second.size());
        for (const auto & access : group.second) {
            shared_ptr<MemBlock> mem(new MemBlock(access->addr, access->addr + access->size, access->timeStart, access->timeDuration));
            mem->record = access->id;
            mem->needRecord = true;
            task->timeStart = mem->timeStart;
            if (mem->timeDuration > task->timeLen) {
                task->timeLen = mem->timeDuration;
            }
            task->memSet.insert(mem);
        }
        taskVec.push_back(task);
    }
    // 生成内存保留链表
    unsigned taskNum = taskVec.size();
    /* set<shared_ptr<MemBlock>> * nextTaskMemVec = NULL; */
    map<unsigned long long, pair<shared_ptr<MemBlock>, unsigned long long>> nextMemMap;
    for (int taskIdx = taskNum - 1; taskIdx >= 0; --taskIdx) {
        if (!nextMemMap.empty()) {
            auto & curTask = taskVec[taskIdx];
            for (auto & mem : curTask->memSet) {
                auto relatedBeginIt = nextMemMap.lower_bound(mem->addrBegin);
                auto relatedEndIt = nextMemMap.upper_bound(mem->addrEnd);
                for (auto it = relatedBeginIt; it != relatedEndIt; ++it) {
                    auto nextMem = it->second.first;
                    if (nextMem->addrBegin >= mem->addrEnd) {
                        continue;
                    }
                    LOG("mem %llu-%llu(%p) -> mem %llu-%llu(%p)", mem->addrBegin, mem->addrEnd, mem.get(), nextMem->addrBegin, nextMem->addrEnd, nextMem.get());
                    mem->nextUseBlock.push_back(nextMem);
                }
            }
        }
        for (auto & mem : taskVec[taskIdx]->memSet) {
            auto beginIt = nextMemMap.lower_bound(mem->addrBegin);
            auto endIt = nextMemMap.lower_bound(mem->addrEnd);
            auto removeBeginIt = beginIt;
            auto removeEndIt = endIt;
            bool modifyEnd = false;
            // 判断是否需要保留beginIt
            if (beginIt != nextMemMap.end()) {
                if (removeEndIt != nextMemMap.end() && mem->addrBegin > beginIt->first + beginIt->second.second) {
                    beginIt->second.second = mem->addrBegin - beginIt->first;
                    ++removeBeginIt;
                }
                else if (mem->addrBegin > beginIt->first) {
                    beginIt->second.second = mem->addrBegin - beginIt->first;
                }
            }
            // 判断是否需要保留endIt
            if (endIt != nextMemMap.end()) {
                if (mem->addrEnd > endIt->first) {
                    if (mem->addrEnd >= endIt->first + endIt->second.second) {
                        ++removeEndIt;
                    }
                    else if (mem->addrEnd > endIt->first) {
                        modifyEnd = true;
                    }
                }
            }
            nextMemMap[mem->addrBegin] = make_pair(mem, mem->addrEnd - mem->addrBegin);
            nextMemMap.erase(removeBeginIt, removeEndIt);
            if (modifyEnd) {
                auto k = endIt->first;
                auto v = endIt->second;
                nextMemMap.erase(endIt);
                v.second = k + v.second - mem->addrEnd;
                k = mem->addrEnd;
                nextMemMap[k] = v;
            }
        }
    }
}

// 主调度循环
void Scheduler::run()
{
    // 按任务的开始时间顺序处理
    unsigned taskNum = taskVec.size();
    LOG("task num %d", taskNum);

    // 缓存内存的下一次使用时间
    vector<shared_ptr<Interval>> waitReleaseIntervalVec;
    map<unsigned long long, vector<Interval *>> usetimeIntervalMap; // <time, interval>
    for (unsigned taskIdx = 0; taskIdx < taskNum; ++taskIdx) {
        waitReleaseIntervalVec.clear();
        usetimeIntervalMap.clear();
        LOG("%d", taskIdx);
        IntervalTree taskMemoryTree;
        prepareTask(taskVec[taskIdx], taskMemoryTree);
        waitReleaseIntervalVec.clear();
        usetimeIntervalMap.clear();
        // 以下两个函数不直接操作currentTime，只返回从现在起的耗时
        // 先获取处理任务的时间，再依次获取加载内存的时间
        unsigned long long preloadNeedTime = 0;
        unsigned long long visitNeedTime = 0;
        unsigned curTaskIdx = taskIdx;
        visitNeedTime = processTask(currentTime, taskVec[taskIdx]);
        unsigned long long totalPreloadNeedTime = 0;
        if (visitNeedTime > 0) { // 考虑time可能为0的情况
            /* 向后直到找到一个不满足的Task */
            for (unsigned i = curTaskIdx + 1; i < taskNum && i < curTaskIdx + 10; ++i) {
                /* 向处理中传入中间所有任务并集的内存需求IntervalTree，提高性能 */
                preloadNeedTime = checkAndPrepareNextTask(currentTime + totalPreloadNeedTime, taskVec[i], taskMemoryTree, visitNeedTime - totalPreloadNeedTime);
                totalPreloadNeedTime += preloadNeedTime;
                LOG("preloadNeedTime: %llu, total: %llu, visit: %llu", preloadNeedTime, totalPreloadNeedTime, visitNeedTime);
                if (totalPreloadNeedTime >= visitNeedTime) {
                    break;
                }
            }
        }
        currentTime += max(totalPreloadNeedTime, visitNeedTime);
    }
    /* 根据两者时间差异进行处理 */
    LOG("Fin %d", currentTime);
    resultOpVec.push_back({"Fin", currentTime, 0, 0});
}

void IntervalTreeCombineMemSet(IntervalTree & tree, set<shared_ptr<MemBlock>> & memSet)
{
    LOG("memSet size %d", memSet.size());
    for (auto & mem : memSet) {
        LOG("%llu, %llu", mem->addrBegin, mem->addrEnd);
        tree.insert(Interval(mem->addrBegin, mem->addrEnd, mem.get()));
    }
}

static vector<Interval> offloadIntervalVec;
static vector<Interval> reloadIntervalVec;

void Scheduler::prepareTask(shared_ptr<Task> & task, IntervalTree & taskMemoryTree)
{
    /* 获取期望的状态 */
    IntervalTreeCombineMemSet(taskMemoryTree, task->memSet);

    memoryTree.print("memory record: ");
    taskMemoryTree.print("expect record: ");

    /* 计算需要加载的内存 */
    memoryTree.getDifference(taskMemoryTree, offloadIntervalVec, reloadIntervalVec);
    LOG("offload %d reload %d", offloadIntervalVec.size(), reloadIntervalVec.size());
    unsigned long long loadSz = 0;
    for (auto & interval : reloadIntervalVec) {
        loadSz += interval.end - interval.start;
    }
    /* 如果内存不足，先释放 */
    /* LOG("memory used %lld load %d", memoryTree.getTotalSize(), loadSz); */
    if (memoryTree.getTotalSize() + loadSz > memoryCapacity) {
        unsigned long long lackSz = memoryTree.getTotalSize() + loadSz - memoryCapacity;
        vector<shared_ptr<Interval>> waitReleaseIntervalVec;
        map<unsigned long long, vector<const Interval *>> usetimeIntervalMap; // <time, interval>
        /* 寻找最近不使用的内存 */
        for (auto & interval : offloadIntervalVec) {
            LOG("mem %llu-%llu(%p)", interval.start, interval.end, interval.user);
            /* 如果不再被使用，直接放到最后 */
            if (interval.user == NULL) {
                LOG("record interval [%llu] %llu-%llu", (unsigned long long)-1, interval.start, interval.end);
                usetimeIntervalMap[(unsigned long long)-1].push_back(&interval);
                continue;
            }
            MemBlock * mem = (MemBlock *)interval.user;
            LOG("user %p, empty %d", interval.user, mem->nextUseBlock.empty());
            if (mem->nextUseBlock.empty()) {
                LOG("record interval [%llu] %llu-%llu", (unsigned long long)-1, interval.start, interval.end);
                usetimeIntervalMap[(unsigned long long)-1].push_back(&interval);
                continue;
            }
            // 先获取存在下一次使用的所有块
            IntervalTree tempTree;
            tempTree.insert(Interval(mem->addrBegin, mem->addrEnd));
            for (auto & nextMem : mem->nextUseBlock) {
                LOG("mem %llu - %llu, nextMem: %llu - %llu", mem->addrBegin, mem->addrEnd, nextMem->addrBegin, nextMem->addrEnd);
                shared_ptr<Interval> newInterval(new Interval(min(interval.start, nextMem->addrBegin), max(interval.end, nextMem->addrEnd), &nextMem));
                waitReleaseIntervalVec.push_back(newInterval);
                LOG("record interval [%llu] %llu-%llu", nextMem->timeStart, newInterval->start, newInterval->end);
                usetimeIntervalMap[nextMem->timeStart].push_back(newInterval.get());
                tempTree.remove(*newInterval);
            }
            // 如果没有下一次使用的块，记录到最后
            for (auto & interval : tempTree.getIntervals()) {
                LOG("record interval [%llu] %llu-%llu", (unsigned long long)-1, interval.start, interval.end);
                shared_ptr<Interval> newInterval(new Interval(interval));
                waitReleaseIntervalVec.push_back(newInterval);
                usetimeIntervalMap[(unsigned long long)-1].push_back(newInterval.get());
            }
        }
        for (map<unsigned long long, vector<const Interval *>>::iterator it = usetimeIntervalMap.end(); it != usetimeIntervalMap.begin() && lackSz > 0;) {
            --it;
            for (auto & interval : it->second) {
                LOG("use time %llu, %llu-%llu", it->first, interval->start, interval->end);
                unsigned long long intervalSz = interval->end - interval->start;
                // 不一定是一块完整的，达到最小要求就行
                unsigned long long offloadSz = intervalSz >= lackSz ? lackSz : intervalSz;
                LOG("Offload %llu intervalSz %llu lackSz %llu", offloadSz, intervalSz, lackSz);
                LOG("Offload %llu %llu %llu", currentTime, interval->start, offloadSz);
                resultOpVec.push_back({"Offload", currentTime, interval->start, offloadSz});
                currentTime += offloadSz * 40;
                memoryTree.remove(Interval(interval->start, interval->start + offloadSz));
                lackSz -= offloadSz;
                if (lackSz == 0) {
                    break;
                }
            }
        }
        if (lackSz > 0) {
            taskMemoryTree.print("ERROR: memory insufficient, expect memory more than capacity ");
            exit(-1);
        }
    }
    /* 加载缺少的内存 */
    for (auto & interval : reloadIntervalVec) {
        LOG("Reload %d %llu %llu", currentTime, interval.start, interval.end - interval.start);
        resultOpVec.push_back({"Reload", currentTime, interval.start, interval.end - interval.start});
        currentTime += (interval.end - interval.start) * 40;
        memoryTree.insert(Interval(interval.start, interval.end, interval.user));
    }
}

unsigned long long Scheduler::checkAndPrepareNextTask(
        unsigned long long startTime,
        shared_ptr<Task> & nextTask,
        IntervalTree & taskMemoryTree,
        unsigned long long maxTime)
{
    unsigned long long elapsed = 0;
    /* 获取期望的状态 */
    IntervalTreeCombineMemSet(taskMemoryTree, nextTask->memSet);

    memoryTree.print("memory record: ");
    taskMemoryTree.print("cur & next expect record: ");

    /* 计算需要加载的内存 */
    memoryTree.getDifference(taskMemoryTree, offloadIntervalVec, reloadIntervalVec);
    LOG("offload %d reload %d", offloadIntervalVec.size(), reloadIntervalVec.size());
    unsigned long long loadSz = 0;
    for (auto & interval : reloadIntervalVec) {
        loadSz += interval.end - interval.start;
    }
    /* 如果内存不足，先释放 */
    if (memoryTree.getTotalSize() + loadSz > memoryCapacity) {
        unsigned long long lackSz = memoryTree.getTotalSize() + loadSz - memoryCapacity;
        vector<shared_ptr<Interval>> waitReleaseIntervalVec;
        map<unsigned long long, vector<const Interval *>> usetimeIntervalMap; // <time, interval>

        /* 寻找最近不使用的内存 */
        for (auto & interval : offloadIntervalVec) {
            /* 如果不再被使用，直接放到最后 */
            if (interval.user == NULL) {
                usetimeIntervalMap[(unsigned long long)-1].push_back(&interval);
                continue;
            }
            MemBlock * mem = (MemBlock *)interval.user;
            if (mem->nextUseBlock.empty()) {
                usetimeIntervalMap[(unsigned long long)-1].push_back(&interval);
                continue;
            }
            // 先获取存在下一次使用的所有块
            IntervalTree tempTree;
            for (auto & nextMem : mem->nextUseBlock) {
                LOG("mem %p %llu - %llu, nextMem: %llu - %llu", mem, mem->addrBegin, mem->addrEnd, nextMem->addrBegin, nextMem->addrEnd);
                shared_ptr<Interval> newInterval(new Interval(min(interval.start, nextMem->addrBegin), max(interval.end, nextMem->addrEnd), &nextMem));
                waitReleaseIntervalVec.push_back(newInterval);
                usetimeIntervalMap[nextMem->timeStart].push_back(newInterval.get());
                tempTree.remove(*newInterval);
            }
            // 如果没有下一次使用的块，记录到最后
            for (auto & interval : tempTree.getIntervals()) {
                usetimeIntervalMap[(unsigned long long)-1].push_back(&interval);
            }
        }
        for (map<unsigned long long, vector<const Interval *>>::iterator it = usetimeIntervalMap.end(); it != usetimeIntervalMap.begin() && lackSz > 0;) {
            --it;
            for (auto & interval : it->second) {
                unsigned long long intervalSz = interval->end - interval->start;
                unsigned long long offloadSz = intervalSz >= lackSz ? lackSz : intervalSz;
                LOG("elapsed %d offloadSz %d maxTime %d", elapsed, offloadSz, maxTime);
                /* 访问时间短也可以卸载一部分内存 */
                if (offloadSz * 40 > maxTime) {
                    offloadSz = (maxTime + 39) / 40;
                }
                LOG("Offload %llu %llu %llu", startTime + elapsed, interval->start, offloadSz);
                resultOpVec.push_back({"Offload", startTime + elapsed, interval->start, offloadSz});
                elapsed += offloadSz * 40;
                LOG("elapsed %d offloadSz %d maxTime %d", elapsed, offloadSz, maxTime);
                memoryTree.remove(Interval(interval->start, interval->start + offloadSz));
                // TODO: 不一定是一块完整的，达到最小要求就行
                lackSz -= offloadSz;
                if (lackSz == 0 || elapsed >= maxTime) {
                    break;
                }
            }
        }
        if (lackSz > 0 || elapsed >= maxTime) {
            return elapsed;
        }
    }
    /* 加载缺少的内存 */
    for (auto & interval : reloadIntervalVec) {
        unsigned long long reloadSz = interval.end - interval.start;
        /* 访问时间短也可以加载一部分内存 */
        if (reloadSz * 40 > (maxTime - elapsed)) {
            reloadSz = (maxTime - elapsed + 39) / 40;
        }
        LOG("Reload %d %llu %llu", startTime + elapsed, interval.start, reloadSz);
        resultOpVec.push_back({"Reload", startTime + elapsed, interval.start, reloadSz});
        elapsed += reloadSz * 40;
        memoryTree.insert(Interval(interval.start, interval.start + reloadSz, interval.user));
        if (elapsed >= maxTime) {
            break;
        }
    }
    return elapsed;
}

// 处理单个任务
unsigned long long Scheduler::processTask(unsigned long long startTime, shared_ptr<Task> & task)
{
    LOG("Task %llu - %llu", task->timeStart, task->timeStart + task->timeLen);
    unsigned long long elapsed = 0;
    LOG("mem %d", task->memSet.size());
    for (auto & mem : task->memSet) {
        /* LOG("mem %p needRecord %d", mem.get(), mem->needRecord); */
        if (mem->needRecord) {
            unsigned long long realStartTime = startTime >= mem->timeStart ? startTime : mem->timeStart;
            LOG("%llu %llu", startTime, mem->timeStart);
            LOG("Visit %llu %d", realStartTime, mem->record);
            resultOpVec.push_back({"Visit", realStartTime, mem->record, 0});
            elapsed = realStartTime + mem->timeDuration - startTime;
        }
    }
    return elapsed;
}

const std::vector<Operation> & Scheduler::getResult() const
{
    return resultOpVec;
}
